"""Rfb data type realization.

Can't use `struct` module becous it's bounded to Python's underlying
C libraries. And given libraries are probably architecture-dependent.


The format (unlike `struct` module) is a iterable of strings described
in RFB protocol documentation:
    U8, U16, U32 -- unsigned types with given length in bytes
    S8, S16, S32 -- signed types.

In fact, it would be best to implement this module in C.
but since I want pure Python implementation....
"""

import struct

# since some structure sizes are dependant on current client mode
# (e.g. on current bytes per pixel setting)
# Struct calculator was made an object that binds to client
# during creation

class TypeObject(object):

    _typeLengths = {
        "x":     1, # padding
        "U8":    1,
        "U16":   2,
        "U32":   4,
        "S8":    1,
        "S16":   2,
        "S32":   4,
        "PIXEL_FORMAT": 16,
    }

    _dynamicTypes = {
        "PIXEL":    lambda s: s._client.getBpp(),
    }

    # complex structure definitions
    _structDefs = {
        "PIXEL_FORMAT": (
            ("bits-per-pixel", "U8"),
            ("depth", "U8"),
            ("big-endian-flag", "U8"),
            ("true-colour-flag", "U8"),
            ("red-max", "U16"),
            ("green-max", "U16"),
            ("blue-max", "U16"),
            ("red-shift", "U8"),
            ("green-shift", "U8"),
            ("blue-shift", "U8"),
            # padding
            ("", "x"), ("", "x"), ("", "x"),
        )
    }

    _client = None

    def __init__(self, client):
        self._client = client

    # biggest possible size (used for packing)
    _structElSize = struct.calcsize(">q")

    def getTypeSize(self, name):
        if name in self._typeLengths:
            # if field has constant length
            _rv = self._typeLengths[name]
        else:
            # else assume it has constant length
            _rv = self._dynamicTypes[name]
            _rv = _rv(self)
        return _rv

    def toInt(self, val):
        _size = self._structElSize
        _val = "\x00" * _size + val
        return struct.unpack(">Q", _val[-_size:])[0]

    def pack(self, format, data):
        _structDefs = self._structDefs
        _structElSize = self._structElSize
        _complexFormats = self._structDefs.keys()
        _formatsWithData = [
            _el for _el in format
            if _el != "x"
        ]
        if len(_formatsWithData) != len(data):
            raise Exception("Format and data element count must be same.")
        _primitiveValues = [ (_typ, _val)
            for (_typ, _val) in zip(_formatsWithData, data)
            if _typ not in _complexFormats ]
        _packFmt = "".join(
            "q" if _fmt.startswith("S") else "Q"
            for (_fmt, _val) in _primitiveValues
        )
        _packedData = struct.pack(">" + _packFmt, \
            *[_val for (_typ, _val) in _primitiveValues])
        _rv = ""
        _dataId = -1
        for (_pos, _fmt) in enumerate(format):
            if _fmt == "x":
                _rv += "\x42"
                continue
            _dataId += 1
            if _fmt in _complexFormats:
                _rv += self.packDict(_fmt, data[_dataId])
            _curEl = _packedData[
                _dataId*_structElSize:(_dataId+1)*_structElSize]
            _neededLen = self.getTypeSize(_fmt)
            _removedLen = _structElSize - _neededLen
            if _removedLen and any(map(
                lambda ch: ch not in ("\xff", "\x00"),
                _curEl[:_removedLen]
            )):
                raise Exception(
                    "Too big number (%r) for such type (%s)." % (
                    _curEl, _fmt))
            _rv += _curEl[_removedLen:]
        return _rv

    def packDict(self, type, data):
        _types = []
        _values = []
        for (_name, _type) in self._structDefs[type]:
            _types.append(_type)
            if _type != "x":
                _values.append(data[_name])
        return self.pack(_types, _values)

    def unpack(self, format, data):
        _rv = []
        # data is in big-endian format
        for _fmt in format:
            _signed = _fmt.startswith("S")
            _len = self.getTypeSize(_fmt)
            _bytes = data[:_len]
            data = data[_len:]
            if _fmt == "x":
                continue
            elif _fmt == "PIXEL_FORMAT":
                _rv.append(self.unpackDict(self._structDefs[_fmt], _bytes))
            elif _fmt in ("PIXEL", "CPIXEL"):
                _rv.append(_bytes)
                continue
            _bytes = [ord(_el) for _el in _bytes]
            _minus = False
            if _signed:
                # in Python, logic is faster that byte arithmetics
                if _bytes[0] >= 0x80:
                    _minus = True
                    _bytes = [_el ^ 255 for _el in _bytes]
            _value = reduce(lambda x, y: (x << 8) + y, _bytes)
            if _minus:
                _value = ( _value + 1 ) * -1
            _rv.append(_value)
        assert not data, "All data should be parsed by now."
        return _rv

    def unpackDict(self, elems, data):
        """Unpack data to dictionary.

        `elems` is list of tuples each containing (name, type).

        """
        _names = [_el[0] for _el in elems if _el[1] != "x"]
        _format = [_el[1] for _el in elems]
        _data = self.unpack(_format, data)
        return dict(zip(_names, _data))

    def calcsize(self, format):
        """Return estimated size for given format.

        Format is ether list of strings or tuples
        that countain (name, type) pairs (
        as used by `unpackDict`)

        """
        if all(map(lambda el: isinstance(el, basestring), format)):
            # all elements are strings
            _fmt = format
        else:
            _fmt = [_el[1] for _el in format]
        return sum(map(lambda _el: self.getTypeSize(_el), _fmt))

# vim: set sts=4 sw=4 et :
